import datetime
import time
import matplotlib.pyplot as plt
import numpy as np


class MetricLogger():
    def __init__(self, save_dir):
        self.save_interval = 10
        self.save_log = save_dir / "log"  # 保存日志的路径
        with open(self.save_log, "w") as f:
            f.write(
                f"{'Episode':>8}{'Step':>8}{'Epsilon':>10}{'MeanReward':>15}"
                f"{'MeanEnergy_UAV':>15}{'MeanEnergy_User':>15}{'MeanAoI_UE':>15}{'MeanAoI_DC':>15}"
                #  f"{'MeanLength':>15}{'MeanLoss':>15}{'MeanQValue':>15}"
                f"{'TimeDelta':>15}{'Time':>20}\n"
            )
        self.ep_rewards_plot = save_dir / "reward_plot.jpg"
        self.ep_energy_uav_plot = save_dir / "energy_uav_plot.jpg"
        self.ep_energy_user_plot = save_dir / "energy_user_plot.jpg"
        self.ep_aoi_ue_plot = save_dir / "aoi_ue_plot.jpg"
        self.ep_aoi_dc_plot = save_dir / "aoi_dc_plot.jpg"
        # Initialize episode metrics
        self.curr_ep_reward = 0.0
        self.curr_ep_energy_uav = 0.0
        self.curr_ep_energy_user = 0.0
        self.curr_ep_aoi_ue = 0.0
        self.curr_ep_aoi_dc = 0.0
        # History metrics
        self.ep_rewards = []

        self.ep_lengths = []
        self.ep_avg_losses = []

        self.ep_energy_uav = []
        self.ep_energy_user = []
        self.ep_aoi_ue = []
        self.ep_aoi_dc = []
        # Moving averages, added for every call to record()
        self.moving_avg_ep_rewards = []
        self.moving_avg_ep_lengths = []
        self.moving_avg_ep_avg_losses = []
        self.moving_avg_ep_energy_uav = []
        self.moving_avg_ep_energy_user = []
        self.moving_avg_ep_aoi_ue = []
        self.moving_avg_ep_aoi_dc = []
        # Current episode metric
        self.init_episode()

        # Timing  修改相关性能参数
        self.record_time = time.time()

    def log_step(self, reward, energy_uav, energy_user, aoi_ue, aoi_dc):
        self.curr_ep_reward += reward
        self.curr_ep_energy_uav += energy_uav
        self.curr_ep_energy_user += energy_user
        self.curr_ep_aoi_ue += aoi_ue
        self.curr_ep_aoi_dc += aoi_dc

    def log_episode(self):
        """Mark end of episode"""
        self.ep_rewards.append(self.curr_ep_reward)
        self.ep_energy_uav.append(self.curr_ep_energy_uav)
        self.ep_energy_user.append(self.curr_ep_energy_user)
        self.ep_aoi_ue.append(self.curr_ep_aoi_ue)
        self.ep_aoi_dc.append(self.curr_ep_aoi_dc)
        self.init_episode()

    def init_episode(self):
        self.curr_ep_reward = 0.0
        self.curr_ep_energy_uav = 0.0
        self.curr_ep_energy_user = 0.0
        self.curr_ep_aoi_ue = 0.0
        self.curr_ep_aoi_dc = 0.0

    # def record_nn(self, nn):
    #     with open(str(self.save_log) + str("_nn_specs"), "a") as f:
    #         f.write(str(summary(nn, (1,1,3))))

    #  记录超参数
    def record_initials(self, memory_len, batch_size, exploration_rate_decay, burnin, learn_every, sync_every):
        with open(str(self.save_log) + str("_specs"), "a") as f:
            f.write(
                f"Memory len: {memory_len:15d}\nbatch size: {batch_size:15.3f}"
                f"\nexploration rate decay: {exploration_rate_decay:15.10f}\nburnin: {burnin:15.3f}\nlearn every: {learn_every:15.3f}"
                f"\nsync every: {sync_every:15.3f}\n"
            )

    # 性能监控和可视化
    def record(self, episode, step):
        # 计算最近10个episode的平均值
        mean_ep_reward = np.round(np.mean(self.ep_rewards[-10:]), 3)
        mean_energy_uav = np.round(np.mean(self.ep_energy_uav[-10:]), 3)
        mean_energy_user = np.round(np.mean(self.ep_energy_user[-10:]), 3)
        mean_aoi_ue = np.round(np.mean(self.ep_aoi_ue[-10:]), 3)
        mean_aoi_dc = np.round(np.mean(self.ep_aoi_dc[-10:]), 3)

        # 更新移动平均列表
        self.moving_avg_ep_rewards.append(mean_ep_reward)
        self.moving_avg_ep_energy_uav.append(mean_energy_uav)
        self.moving_avg_ep_energy_user.append(mean_energy_user)
        self.moving_avg_ep_aoi_ue.append(mean_aoi_ue)
        self.moving_avg_ep_aoi_dc.append(mean_aoi_dc)

        # 记录当前时间
        self.record_time = time.time()

        # 输出每个指标的值
        print(
            f"Episode {episode} - "
            f"Step {step} - "
            f"Mean Reward: {mean_ep_reward:.3f} - "
            f"Mean Energy UAV: {mean_energy_uav:.3f} - "
            f"Mean Energy User: {mean_energy_user:.3f} - "
            f"Mean AoI UE: {mean_aoi_ue:.3f} - "
            f"Mean AoI DC: {mean_aoi_dc:.3f} - "
            f"Time: {datetime.datetime.now().strftime('%Y-%m-%dT%H:%M:%S')}"
        )

        # 记录到日志文件
        with open(self.save_log, "a") as f:
            f.write(
                f"{episode:8d}{step:8d}"
                f"{mean_ep_reward:15.3f}{mean_energy_uav:15.3f}{mean_energy_user:15.3f}"
                f"{mean_aoi_ue:15.3f}{mean_aoi_dc:15.3f}"
                f"{datetime.datetime.now().strftime('%Y-%m-%dT%H:%M:%S'):>20}\n"
            )

        # 绘制并保存每个指标的图形
        for metric in ["ep_rewards", "ep_energy_uav", "ep_energy_user", "ep_aoi_ue", "ep_aoi_dc"]:
            plt.figure(figsize=(10, 5))
            plt.plot(getattr(self, f"moving_avg_{metric}"), label=f"Moving Avg {metric}")
            plt.title(f"Moving Average of {metric.replace('_', ' ').title()}")
            plt.xlabel('Episodes')
            plt.ylabel(metric.replace('_', ' ').title())
            plt.grid()
            plt.legend()
            plt.savefig(getattr(self, f"{metric}_plot"))
            plt.clf()  # 清空图形以便绘制下一个

